import os
import sys
import cv2
import torch
from PyQt5.QtWidgets import (QApplication, QMainWindow, QPushButton, QVBoxLayout,
                             QWidget, QLabel, QListWidget, QScrollArea, QGridLayout,
                             QMessageBox, QAbstractItemView, QHBoxLayout, QSizePolicy)
from PyQt5.QtGui import QImage, QPixmap, QFont, QPainter, QColor
from PyQt5.QtCore import QTimer, QDateTime, Qt


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("FOC Detection Pro")
        self.setGeometry(100, 100, 1600, 900)

        # 主界面布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)

        # 初始化各组件
        self.setup_selection_panel()
        self.setup_video_display()
        self.setup_status_bar()

        # 设置布局拉伸系数
        self.main_layout.setStretch(0, 0)  # 选择面板
        self.main_layout.setStretch(1, 1)  # 视频区域
        self.main_layout.setStretch(2, 0)  # 状态栏

        # 初始化组件
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frames)
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', _verbose=False)
        self.detect_cameras()

        # 状态变量
        self.showing_stats = False
        self.start_time = None
        self.camera_data = {}
        self.recording = False

        # 添加样式表（关键修改）
        self.setStyleSheet("""
            /* 主界面标签样式 */
            QMainWindow QLabel {
                background: #000;
                color: white;
            }
            /* 网格布局边距 */
            QGridLayout {
                margin: 5px;
            }
            /* 消息框专用样式 */
            QMessageBox {
                background-color: white;
            }
            QMessageBox QLabel {
                color: black;
                background: transparent;
            }
            QMessageBox QPushButton {
                min-width: 80px;
            }
        """)

    def setup_selection_panel(self):
        """摄像头选择面板"""
        self.selection_panel = QWidget()
        panel_layout = QVBoxLayout(self.selection_panel)

        # 摄像头列表
        self.camera_list = QListWidget()
        self.camera_list.setSelectionMode(QAbstractItemView.MultiSelection)
        self.camera_list.setFixedHeight(200)
        panel_layout.addWidget(QLabel("选择摄像头（可多选）："))
        panel_layout.addWidget(self.camera_list)

        # 控制按钮
        btn_layout = QHBoxLayout()
        self.confirm_btn = QPushButton("开始检测", clicked=self.start_detection)
        self.refresh_btn = QPushButton("刷新列表", clicked=self.detect_cameras)
        btn_layout.addWidget(self.refresh_btn)
        btn_layout.addWidget(self.confirm_btn)
        panel_layout.addLayout(btn_layout)

        self.main_layout.addWidget(self.selection_panel)

    def setup_video_display(self):
        """视频显示区域"""
        self.video_container = QWidget()
        self.video_layout = QGridLayout(self.video_container)
        self.video_layout.setContentsMargins(10, 10, 10, 10)

        # 设置布局策略
        self.video_container.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.video_layout.setRowStretch(0, 1)
        self.video_layout.setRowStretch(1, 1)
        self.video_layout.setColumnStretch(0, 1)
        self.video_layout.setColumnStretch(1, 1)

        # 直接添加视频容器到主布局
        self.main_layout.addWidget(self.video_container, 1)

        # 添加记录控制按钮
        self.setup_record_controls()

    def setup_status_bar(self):
        """状态信息栏"""
        self.status_bar = QLabel()
        self.status_bar.setAlignment(Qt.AlignCenter)
        # 移除此处的通用样式设置
        self.status_bar.setStyleSheet("""
            font-size: 14px; 
            color: #333; 
            padding: 5px;
            background: transparent;  /* 新增透明背景 */
        """)
        self.main_layout.addWidget(self.status_bar)

    def detect_cameras(self):
        """检测可用摄像头"""
        self.camera_list.clear()
        for i in range(10):
            cap = cv2.VideoCapture(i)
            if cap.isOpened():
                self.camera_list.addItem(f"Camera {i}")
                cap.release()

    def start_detection(self):
        """开始检测（带异常处理）"""
        try:
            selected = self.camera_list.selectedItems()
            if not selected:
                QMessageBox.warning(self, "警告", "请选择至少一个摄像头")
                return

            # 隐藏选择面板
            self.selection_panel.hide()
            self.video_container.show()

            # 显示加载提示
            loading = QMessageBox(self)
            loading.setWindowTitle("系统提示")
            loading.setText("摄像头初始化中...")
            loading.setStandardButtons(QMessageBox.NoButton)
            loading.show()

            # 异步初始化
            QTimer.singleShot(100, lambda: self._finish_detection(selected, loading))
        except Exception as e:
            QMessageBox.critical(self, "错误", f"初始化失败：{str(e)}")

    def _finish_detection(self, selected, loading):
        """完成检测初始化"""
        try:
            self.init_cameras(selected)
            self.timer.start(30)
            self.start_time = QDateTime.currentDateTime()
        finally:
            loading.hide()

    def init_cameras(self, selected):
        """初始化选中的摄像头"""
        # 清理已有资源
        if hasattr(self, 'caps'):
            for cap in self.caps:
                cap.release()
            for label in self.labels:
                label.deleteLater()

        # 创建新实例
        self.caps = []
        self.labels = []
        self.camera_data.clear()

        # 动态计算布局
        cam_count = len(selected)
        cols = 2 if cam_count > 2 else 1
        rows = (cam_count + cols - 1) // cols

        # 彻底清理旧布局（关键修改）
        while self.video_layout.count():
            item = self.video_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()

        # 重置行列设置
        for r in range(self.video_layout.rowCount()):
            self.video_layout.setRowStretch(r, 0)
        for c in range(self.video_layout.columnCount()):
            self.video_layout.setColumnStretch(c, 0)

        # 设置新行列
        for r in range(rows):
            self.video_layout.setRowStretch(r, 1)
        for c in range(cols):
            self.video_layout.setColumnStretch(c, 1)

        # 添加摄像头
        for i, item in enumerate(selected):
            index = int(item.text().split()[-1])
            cap = cv2.VideoCapture(index)
            if not cap.isOpened():
                continue

            # 初始化摄像头数据
            cam_id = f"camera_{index}"
            self.camera_data[cam_id] = {
                'foc_history': [],
                'current_percent': 0.0,
                'frame_size': None
            }

            # 创建显示标签
            label = VideoLabel(cam_id)
            label.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Ignored)
            label.setAlignment(Qt.AlignCenter)
            label.setStyleSheet("border: 2px solid #666; background: #000;")

            row = i // cols
            col = i % cols
            self.video_layout.addWidget(label, row, col)

            self.caps.append(cap)
            self.labels.append(label)

        # 强制刷新布局
        self.video_container.updateGeometry()
        self.adjustSize()

    def update_frames(self):
        """更新所有摄像头画面"""
        for idx, (cap, label) in enumerate(zip(self.caps, self.labels)):
            ret, frame = cap.read()
            if ret:
                # 模型推理
                results = self.model(frame)

                # 分析结果
                cam_id = label.cam_id
                detections = self.analyze_results(results)
                self.update_camera_data(cam_id, detections)

                # 可视化处理
                frame = self.visualize_frame(frame, results, detections, cam_id)

                # 显示画面
                self.display_frame(frame, label)

        # 更新状态信息
        self.update_status()

    def analyze_results(self, results):
        """分析检测结果"""
        foc_count = sum(1 for x in results.xyxy[0] if self.model.names[int(x[5])] == "foc")
        no_count = sum(1 for x in results.xyxy[0] if self.model.names[int(x[5])] == "no")
        total = foc_count + no_count
        return {
            'foc_percent': (foc_count / total * 100) if total > 0 else 0,
            'foc_count': foc_count,
            'no_count': no_count
        }

    def update_camera_data(self, cam_id, detections):
        """更新摄像头数据"""
        data = self.camera_data[cam_id]
        data['current_percent'] = detections['foc_percent']

        # 仅在记录状态时保存历史数据
        if self.recording:
            if 'recording_history' not in data:
                data['recording_history'] = []
            data['recording_history'].append(detections['foc_percent'])

    def visualize_frame(self, frame, results, detections, cam_id):
        """可视化处理帧"""
        # 边框颜色逻辑
        percent = detections['foc_percent']
        red = 255 if percent < 50 else int(255 * (1 - (percent - 50) / 50))
        green = int(255 * (percent / 50)) if percent < 50 else 255

        # 闪烁效果
        border_color = (0, green, red)
        if percent < 0.1:
            border_color = (0, 0, 0) if (QDateTime.currentMSecsSinceEpoch() // 500) % 2 else border_color

        # 绘制处理
        cv2.rectangle(frame, (0, 0), (frame.shape[1], frame.shape[0]), border_color, 15)
        frame = cv2.cvtColor(results.render()[0], cv2.COLOR_BGR2RGB)

        # 添加文字信息
        text = f"FOC: {percent:.1f}% ({detections['foc_count']}/{detections['no_count']})"
        cv2.putText(frame, text, (20, 50), cv2.FONT_HERSHEY_SIMPLEX,
                    1.2, (255, 255, 255), 3, cv2.LINE_AA)

        return frame

    def display_frame(self, frame, label):
        """显示帧到QLabel"""
        h, w, ch = frame.shape
        bytes_per_line = ch * w
        q_img = QImage(frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_img)

        # 使用availableSize代替实际尺寸（关键修改）
        available_size = label.size().boundedTo(pixmap.size())
        scaled_pixmap = pixmap.scaled(
            available_size,
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        label.setPixmap(scaled_pixmap)

    def update_status(self):
        """更新状态信息"""
        if self.camera_data:
            status_parts = []
            total_percent = 0.0
            for cam_id, data in self.camera_data.items():
                status_parts.append(f"{cam_id}: {data['current_percent']:.1f}%")
                total_percent += data['current_percent']

            avg_percent = total_percent / len(self.camera_data)
            duration = self.start_time.secsTo(QDateTime.currentDateTime())
            time_info = f"运行时间: {duration // 60:02d}:{duration % 60:02d}"

            if self.showing_stats:
                stats_info = f" | 平均FOC: {avg_percent:.1f}%"
                self.status_bar.setText(f"{time_info}{stats_info} | {' | '.join(status_parts)}")
            else:
                self.status_bar.setText(f"{time_info} | {' | '.join(status_parts)}")

    def toggle_stats(self):
        """切换统计显示"""
        self.showing_stats = not self.showing_stats
        if not self.showing_stats:
            self.save_statistics()
        self.update_status()

    def setup_record_controls(self):
        """在视频显示区域添加记录控制按钮"""
        self.record_btn = QPushButton("开始记录", clicked=self.toggle_recording)
        self.record_btn.setFixedSize(120, 40)
        self.record_btn.setStyleSheet("""
            QPushButton {
                background: #4CAF50;
                color: white;
                border-radius: 5px;
                font-size: 14px;
            }
            QPushButton:pressed {
                background: #45a049;
            }
        """)

        # 创建控制面板容器
        control_panel = QWidget()
        panel_layout = QHBoxLayout(control_panel)
        panel_layout.addWidget(self.record_btn)
        panel_layout.addStretch()

        # 将控制面板添加到主布局
        self.main_layout.insertWidget(2, control_panel)

    def toggle_recording(self):
        """切换记录状态"""
        self.recording = not self.recording
        if self.recording:
            # 初始化记录状态
            self.record_start_time = QDateTime.currentDateTime()
            for cam_data in self.camera_data.values():
                cam_data['recording_history'] = []
            self.record_btn.setText("停止记录")
            self.record_btn.setStyleSheet("background: #f44336; color: white;")
            QMessageBox.information(self, "记录状态", "开始记录检测数据")
        else:
            # 保存并重置记录数据
            self.save_recording_data()
            self.record_btn.setText("开始记录")
            self.record_btn.setStyleSheet("background: #4CAF50; color: white;")
            QMessageBox.information(self, "记录状态", "数据已保存并重置")

    def save_recording_data(self):
        """保存记录数据到recordings文件夹"""
        if not hasattr(self, 'record_start_time'):
            return

        # 创建存储目录
        save_dir = "recordings"
        os.makedirs(save_dir, exist_ok=True)  # 自动创建目录

        # 生成带时间戳的文件名
        timestamp = self.record_start_time.toString("yyyyMMdd_hhmmss")
        filename = os.path.join(save_dir, f"{timestamp}.txt")

        stats = [
            f"=== 记录时段 {self.record_start_time.toString('yyyy-MM-dd hh:mm:ss')}"
            f" 至 {QDateTime.currentDateTime().toString('yyyy-MM-dd hh:mm:ss')} ==="
        ]

        try:
            for cam_id, data in self.camera_data.items():
                if 'recording_history' in data and data['recording_history']:
                    # 计算统计指标
                    history = data['recording_history']
                    avg = sum(history) / len(history)
                    maximum = max(history)
                    minimum = min(history)

                    # 格式化摄像头数据
                    cam_stats = (
                        f"[{cam_id}]\n"
                        f"平均比例: {avg:.2f}%\n"
                        f"最高比例: {maximum:.2f}%\n"
                        f"最低比例: {minimum:.2f}%\n"
                        f"数据点数: {len(history)}\n"
                    )
                    stats.append(cam_stats)
                    data['recording_history'] = []  # 清空记录数据

            # 写入文件
            with open(filename, "w", encoding="utf-8") as f:
                f.write("\n".join(stats))

            # 显示保存路径提示
            QMessageBox.information(
                self,
                "保存成功",
                f"数据已保存至：\n{os.path.abspath(filename)}",
                QMessageBox.Ok
            )

        except PermissionError:
            QMessageBox.critical(self, "错误",
                                 "文件写入权限被拒绝，请尝试以下操作：\n1. 关闭已打开的文件\n2. 以管理员身份运行程序")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存失败：{str(e)}")

    def closeEvent(self, event):
        """关闭程序时的清理操作"""
        if hasattr(self, 'caps'):
            for cap in self.caps:
                cap.release()
        event.accept()


class VideoLabel(QLabel):  # ← 这里就是定义
    """自定义视频显示组件"""

    def __init__(self, cam_id, parent=None):
        super().__init__(parent)
        self.cam_id = cam_id
        self.setFont(QFont("Arial", 10))

    def paintEvent(self, event):
        """自定义绘制事件"""
        super().paintEvent(event)
        if self.pixmap():
            painter = QPainter(self)
            painter.setPen(QColor(255, 255, 255))
            painter.drawText(10, 30, f"Camera: {self.cam_id}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
